{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "from torch import nn\n",
    "from torch import optim\n",
    "import torch.nn.functional as F\n",
    "from torch.autograd import Variable\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import transforms\n",
    "from torchvision.utils import save_image\n",
    "from torchvision.datasets import MNIST\n",
    "import os\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists('./vae_img'):\n",
    "    os.mkdir('./vae_img')\n",
    "\n",
    "\n",
    "def to_img(x):\n",
    "    x = x.clamp(0, 1)\n",
    "    x = x.view(x.size(0), 1, 28, 28)\n",
    "    return x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VAE(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(VAE, self).__init__()\n",
    "\n",
    "        self.fc1 = nn.Linear(784, 400)\n",
    "        self.fc21 = nn.Linear(400, 20)\n",
    "        self.fc22 = nn.Linear(400, 20)\n",
    "        self.fc3 = nn.Linear(20, 400)\n",
    "        self.fc4 = nn.Linear(400, 784)\n",
    "\n",
    "    def encode(self, x):\n",
    "        h1 = F.relu(self.fc1(x))\n",
    "        return self.fc21(h1), self.fc22(h1)\n",
    "\n",
    "    def reparametrize(self, mu, logvar):\n",
    "        std = logvar.mul(0.5).exp_()\n",
    "        if torch.cuda.is_available():\n",
    "            eps = torch.cuda.FloatTensor(std.size()).normal_()\n",
    "        else:\n",
    "            eps = torch.FloatTensor(std.size()).normal_()\n",
    "        eps = Variable(eps)\n",
    "        return eps.mul(std).add_(mu)\n",
    "\n",
    "    def decode(self, z):\n",
    "        h3 = F.relu(self.fc3(z))\n",
    "        return F.sigmoid(self.fc4(h3))\n",
    "\n",
    "    def forward(self, x):\n",
    "        mu, logvar = self.encode(x)\n",
    "        z = self.reparametrize(mu, logvar)\n",
    "        return self.decode(z), mu, logvar\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_function(recon_x, x, mu, logvar):\n",
    "    \"\"\"\n",
    "    recon_x: generating images\n",
    "    x: origin images\n",
    "    mu: latent mean\n",
    "    logvar: latent log variance\n",
    "    \"\"\"\n",
    "    BCE = reconstruction_function(recon_x, x)  # mse loss\n",
    "    # loss = 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)\n",
    "    KLD_element = mu.pow(2).add_(logvar.exp()).mul_(-1).add_(1).add_(logvar)\n",
    "    KLD = torch.sum(KLD_element).mul_(-0.5)\n",
    "    # KL divergence\n",
    "    return BCE + KLD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz\n",
      "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz\n",
      "Processing...\n",
      "Done!\n",
      "Train Epoch: 0 [0/60000 (0%)]\tLoss: 183.709747\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python2.7/site-packages/ipykernel_launcher.py:34: UserWarning: invalid index of a 0-dim tensor. This will be an error in PyTorch 0.5. Use tensor.item() to convert a 0-dim tensor to a Python number\n",
      "/usr/local/lib/python2.7/site-packages/ipykernel_launcher.py:41: UserWarning: invalid index of a 0-dim tensor. This will be an error in PyTorch 0.5. Use tensor.item() to convert a 0-dim tensor to a Python number\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Epoch: 0 [12800/60000 (21%)]\tLoss: 51.337727\n",
      "Train Epoch: 0 [25600/60000 (43%)]\tLoss: 43.495949\n",
      "Train Epoch: 0 [38400/60000 (64%)]\tLoss: 39.463840\n",
      "Train Epoch: 0 [51200/60000 (85%)]\tLoss: 38.365051\n",
      "====> Epoch: 0 Average loss: 45.5803\n",
      "Train Epoch: 1 [0/60000 (0%)]\tLoss: 36.135548\n",
      "Train Epoch: 1 [12800/60000 (21%)]\tLoss: 34.876072\n",
      "Train Epoch: 1 [25600/60000 (43%)]\tLoss: 33.696060\n",
      "Train Epoch: 1 [38400/60000 (64%)]\tLoss: 33.892876\n",
      "Train Epoch: 1 [51200/60000 (85%)]\tLoss: 33.307556\n",
      "====> Epoch: 1 Average loss: 35.1271\n",
      "Train Epoch: 2 [0/60000 (0%)]\tLoss: 33.927097\n",
      "Train Epoch: 2 [12800/60000 (21%)]\tLoss: 33.718094\n",
      "Train Epoch: 2 [25600/60000 (43%)]\tLoss: 33.017609\n",
      "Train Epoch: 2 [38400/60000 (64%)]\tLoss: 33.364990\n",
      "Train Epoch: 2 [51200/60000 (85%)]\tLoss: 32.442837\n",
      "====> Epoch: 2 Average loss: 33.2538\n",
      "Train Epoch: 3 [0/60000 (0%)]\tLoss: 31.611500\n",
      "Train Epoch: 3 [12800/60000 (21%)]\tLoss: 32.095543\n",
      "Train Epoch: 3 [25600/60000 (43%)]\tLoss: 31.691576\n",
      "Train Epoch: 3 [38400/60000 (64%)]\tLoss: 32.574341\n",
      "Train Epoch: 3 [51200/60000 (85%)]\tLoss: 31.296017\n",
      "====> Epoch: 3 Average loss: 32.3376\n",
      "Train Epoch: 4 [0/60000 (0%)]\tLoss: 31.463570\n",
      "Train Epoch: 4 [12800/60000 (21%)]\tLoss: 31.262827\n",
      "Train Epoch: 4 [25600/60000 (43%)]\tLoss: 31.659679\n",
      "Train Epoch: 4 [38400/60000 (64%)]\tLoss: 29.896202\n",
      "Train Epoch: 4 [51200/60000 (85%)]\tLoss: 33.093983\n",
      "====> Epoch: 4 Average loss: 31.8081\n",
      "Train Epoch: 5 [0/60000 (0%)]\tLoss: 30.515318\n",
      "Train Epoch: 5 [12800/60000 (21%)]\tLoss: 32.035904\n",
      "Train Epoch: 5 [25600/60000 (43%)]\tLoss: 31.026447\n",
      "Train Epoch: 5 [38400/60000 (64%)]\tLoss: 32.154472\n",
      "Train Epoch: 5 [51200/60000 (85%)]\tLoss: 32.589756\n",
      "====> Epoch: 5 Average loss: 31.4348\n",
      "Train Epoch: 6 [0/60000 (0%)]\tLoss: 30.265038\n",
      "Train Epoch: 6 [12800/60000 (21%)]\tLoss: 31.786367\n",
      "Train Epoch: 6 [25600/60000 (43%)]\tLoss: 30.729908\n",
      "Train Epoch: 6 [38400/60000 (64%)]\tLoss: 31.773233\n",
      "Train Epoch: 6 [51200/60000 (85%)]\tLoss: 31.414577\n",
      "====> Epoch: 6 Average loss: 31.1736\n",
      "Train Epoch: 7 [0/60000 (0%)]\tLoss: 31.110806\n",
      "Train Epoch: 7 [12800/60000 (21%)]\tLoss: 30.713400\n",
      "Train Epoch: 7 [25600/60000 (43%)]\tLoss: 32.213604\n",
      "Train Epoch: 7 [38400/60000 (64%)]\tLoss: 28.416868\n",
      "Train Epoch: 7 [51200/60000 (85%)]\tLoss: 30.741480\n",
      "====> Epoch: 7 Average loss: 30.9627\n",
      "Train Epoch: 8 [0/60000 (0%)]\tLoss: 28.803867\n",
      "Train Epoch: 8 [12800/60000 (21%)]\tLoss: 30.302649\n",
      "Train Epoch: 8 [25600/60000 (43%)]\tLoss: 30.214563\n",
      "Train Epoch: 8 [38400/60000 (64%)]\tLoss: 31.026213\n",
      "Train Epoch: 8 [51200/60000 (85%)]\tLoss: 30.352478\n",
      "====> Epoch: 8 Average loss: 30.7836\n",
      "Train Epoch: 9 [0/60000 (0%)]\tLoss: 31.166786\n",
      "Train Epoch: 9 [12800/60000 (21%)]\tLoss: 30.245884\n",
      "Train Epoch: 9 [25600/60000 (43%)]\tLoss: 30.966717\n",
      "Train Epoch: 9 [38400/60000 (64%)]\tLoss: 29.632961\n",
      "Train Epoch: 9 [51200/60000 (85%)]\tLoss: 30.163986\n",
      "====> Epoch: 9 Average loss: 30.6282\n",
      "Train Epoch: 10 [0/60000 (0%)]\tLoss: 30.757961\n",
      "Train Epoch: 10 [12800/60000 (21%)]\tLoss: 31.491070\n",
      "Train Epoch: 10 [25600/60000 (43%)]\tLoss: 30.462280\n",
      "Train Epoch: 10 [38400/60000 (64%)]\tLoss: 30.283237\n",
      "Train Epoch: 10 [51200/60000 (85%)]\tLoss: 30.162361\n",
      "====> Epoch: 10 Average loss: 30.4994\n",
      "Train Epoch: 11 [0/60000 (0%)]\tLoss: 30.466303\n",
      "Train Epoch: 11 [12800/60000 (21%)]\tLoss: 29.423326\n",
      "Train Epoch: 11 [25600/60000 (43%)]\tLoss: 30.197920\n",
      "Train Epoch: 11 [38400/60000 (64%)]\tLoss: 30.531265\n",
      "Train Epoch: 11 [51200/60000 (85%)]\tLoss: 29.630005\n",
      "====> Epoch: 11 Average loss: 30.4079\n",
      "Train Epoch: 12 [0/60000 (0%)]\tLoss: 30.846107\n",
      "Train Epoch: 12 [12800/60000 (21%)]\tLoss: 30.612965\n",
      "Train Epoch: 12 [25600/60000 (43%)]\tLoss: 31.095856\n",
      "Train Epoch: 12 [38400/60000 (64%)]\tLoss: 30.568771\n",
      "Train Epoch: 12 [51200/60000 (85%)]\tLoss: 30.976185\n",
      "====> Epoch: 12 Average loss: 30.3056\n",
      "Train Epoch: 13 [0/60000 (0%)]\tLoss: 29.186455\n",
      "Train Epoch: 13 [12800/60000 (21%)]\tLoss: 29.906204\n",
      "Train Epoch: 13 [25600/60000 (43%)]\tLoss: 29.879099\n",
      "Train Epoch: 13 [38400/60000 (64%)]\tLoss: 31.018965\n",
      "Train Epoch: 13 [51200/60000 (85%)]\tLoss: 29.830441\n",
      "====> Epoch: 13 Average loss: 30.1919\n",
      "Train Epoch: 14 [0/60000 (0%)]\tLoss: 29.806835\n",
      "Train Epoch: 14 [12800/60000 (21%)]\tLoss: 28.663317\n",
      "Train Epoch: 14 [25600/60000 (43%)]\tLoss: 28.556297\n",
      "Train Epoch: 14 [38400/60000 (64%)]\tLoss: 31.297682\n",
      "Train Epoch: 14 [51200/60000 (85%)]\tLoss: 30.525665\n",
      "====> Epoch: 14 Average loss: 30.1131\n",
      "Train Epoch: 15 [0/60000 (0%)]\tLoss: 30.350090\n",
      "Train Epoch: 15 [12800/60000 (21%)]\tLoss: 29.815151\n",
      "Train Epoch: 15 [25600/60000 (43%)]\tLoss: 28.814440\n",
      "Train Epoch: 15 [38400/60000 (64%)]\tLoss: 30.266960\n",
      "Train Epoch: 15 [51200/60000 (85%)]\tLoss: 29.717367\n",
      "====> Epoch: 15 Average loss: 30.0309\n",
      "Train Epoch: 16 [0/60000 (0%)]\tLoss: 29.348808\n",
      "Train Epoch: 16 [12800/60000 (21%)]\tLoss: 29.613543\n",
      "Train Epoch: 16 [25600/60000 (43%)]\tLoss: 30.008385\n",
      "Train Epoch: 16 [38400/60000 (64%)]\tLoss: 29.614532\n",
      "Train Epoch: 16 [51200/60000 (85%)]\tLoss: 30.353191\n",
      "====> Epoch: 16 Average loss: 29.9345\n",
      "Train Epoch: 17 [0/60000 (0%)]\tLoss: 30.188900\n",
      "Train Epoch: 17 [12800/60000 (21%)]\tLoss: 30.463730\n",
      "Train Epoch: 17 [25600/60000 (43%)]\tLoss: 30.326138\n",
      "Train Epoch: 17 [38400/60000 (64%)]\tLoss: 29.962700\n",
      "Train Epoch: 17 [51200/60000 (85%)]\tLoss: 29.509174\n",
      "====> Epoch: 17 Average loss: 29.8562\n",
      "Train Epoch: 18 [0/60000 (0%)]\tLoss: 31.719082\n",
      "Train Epoch: 18 [12800/60000 (21%)]\tLoss: 29.054089\n",
      "Train Epoch: 18 [25600/60000 (43%)]\tLoss: 29.195709\n",
      "Train Epoch: 18 [38400/60000 (64%)]\tLoss: 30.419323\n",
      "Train Epoch: 18 [51200/60000 (85%)]\tLoss: 30.445457\n",
      "====> Epoch: 18 Average loss: 29.8053\n",
      "Train Epoch: 19 [0/60000 (0%)]\tLoss: 29.747269\n",
      "Train Epoch: 19 [12800/60000 (21%)]\tLoss: 29.747875\n",
      "Train Epoch: 19 [25600/60000 (43%)]\tLoss: 27.798092\n",
      "Train Epoch: 19 [38400/60000 (64%)]\tLoss: 30.838982\n",
      "Train Epoch: 19 [51200/60000 (85%)]\tLoss: 29.384138\n",
      "====> Epoch: 19 Average loss: 29.7259\n",
      "Train Epoch: 20 [0/60000 (0%)]\tLoss: 30.144314\n",
      "Train Epoch: 20 [12800/60000 (21%)]\tLoss: 29.369972\n",
      "Train Epoch: 20 [25600/60000 (43%)]\tLoss: 29.801529\n",
      "Train Epoch: 20 [38400/60000 (64%)]\tLoss: 29.754200\n",
      "Train Epoch: 20 [51200/60000 (85%)]\tLoss: 30.010479\n",
      "====> Epoch: 20 Average loss: 29.6844\n",
      "Train Epoch: 21 [0/60000 (0%)]\tLoss: 30.517113\n",
      "Train Epoch: 21 [12800/60000 (21%)]\tLoss: 29.072678\n",
      "Train Epoch: 21 [25600/60000 (43%)]\tLoss: 28.859062\n",
      "Train Epoch: 21 [38400/60000 (64%)]\tLoss: 30.250624\n",
      "Train Epoch: 21 [51200/60000 (85%)]\tLoss: 29.476585\n",
      "====> Epoch: 21 Average loss: 29.6187\n",
      "Train Epoch: 22 [0/60000 (0%)]\tLoss: 30.098328\n",
      "Train Epoch: 22 [12800/60000 (21%)]\tLoss: 28.629471\n",
      "Train Epoch: 22 [25600/60000 (43%)]\tLoss: 30.572609\n",
      "Train Epoch: 22 [38400/60000 (64%)]\tLoss: 28.068541\n",
      "Train Epoch: 22 [51200/60000 (85%)]\tLoss: 29.392380\n",
      "====> Epoch: 22 Average loss: 29.5821\n",
      "Train Epoch: 23 [0/60000 (0%)]\tLoss: 30.435585\n",
      "Train Epoch: 23 [12800/60000 (21%)]\tLoss: 30.050661\n",
      "Train Epoch: 23 [25600/60000 (43%)]\tLoss: 29.320026\n",
      "Train Epoch: 23 [38400/60000 (64%)]\tLoss: 28.943886\n",
      "Train Epoch: 23 [51200/60000 (85%)]\tLoss: 28.879057\n",
      "====> Epoch: 23 Average loss: 29.5243\n",
      "Train Epoch: 24 [0/60000 (0%)]\tLoss: 29.293987\n",
      "Train Epoch: 24 [12800/60000 (21%)]\tLoss: 29.684799\n",
      "Train Epoch: 24 [25600/60000 (43%)]\tLoss: 30.122471\n",
      "Train Epoch: 24 [38400/60000 (64%)]\tLoss: 30.362911\n",
      "Train Epoch: 24 [51200/60000 (85%)]\tLoss: 28.564873\n",
      "====> Epoch: 24 Average loss: 29.4494\n",
      "Train Epoch: 25 [0/60000 (0%)]\tLoss: 29.118896\n",
      "Train Epoch: 25 [12800/60000 (21%)]\tLoss: 29.253654\n",
      "Train Epoch: 25 [25600/60000 (43%)]\tLoss: 29.483219\n",
      "Train Epoch: 25 [38400/60000 (64%)]\tLoss: 30.634323\n",
      "Train Epoch: 25 [51200/60000 (85%)]\tLoss: 27.927185\n",
      "====> Epoch: 25 Average loss: 29.4577\n",
      "Train Epoch: 26 [0/60000 (0%)]\tLoss: 31.087448\n",
      "Train Epoch: 26 [12800/60000 (21%)]\tLoss: 29.190834\n",
      "Train Epoch: 26 [25600/60000 (43%)]\tLoss: 30.075161\n",
      "Train Epoch: 26 [38400/60000 (64%)]\tLoss: 29.303101\n",
      "Train Epoch: 26 [51200/60000 (85%)]\tLoss: 29.212181\n",
      "====> Epoch: 26 Average loss: 29.3775\n",
      "Train Epoch: 27 [0/60000 (0%)]\tLoss: 29.102173\n",
      "Train Epoch: 27 [12800/60000 (21%)]\tLoss: 28.687901\n",
      "Train Epoch: 27 [25600/60000 (43%)]\tLoss: 29.175978\n",
      "Train Epoch: 27 [38400/60000 (64%)]\tLoss: 28.960945\n",
      "Train Epoch: 27 [51200/60000 (85%)]\tLoss: 30.722912\n",
      "====> Epoch: 27 Average loss: 29.3388\n",
      "Train Epoch: 28 [0/60000 (0%)]\tLoss: 29.098724\n",
      "Train Epoch: 28 [12800/60000 (21%)]\tLoss: 29.216530\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Epoch: 28 [25600/60000 (43%)]\tLoss: 30.105228\n",
      "Train Epoch: 28 [38400/60000 (64%)]\tLoss: 28.444664\n",
      "Train Epoch: 28 [51200/60000 (85%)]\tLoss: 28.852636\n",
      "====> Epoch: 28 Average loss: 29.3043\n",
      "Train Epoch: 29 [0/60000 (0%)]\tLoss: 29.288792\n",
      "Train Epoch: 29 [12800/60000 (21%)]\tLoss: 29.583843\n",
      "Train Epoch: 29 [25600/60000 (43%)]\tLoss: 30.936615\n",
      "Train Epoch: 29 [38400/60000 (64%)]\tLoss: 29.512413\n",
      "Train Epoch: 29 [51200/60000 (85%)]\tLoss: 29.605236\n",
      "====> Epoch: 29 Average loss: 29.2677\n",
      "Train Epoch: 30 [0/60000 (0%)]\tLoss: 28.870108\n",
      "Train Epoch: 30 [12800/60000 (21%)]\tLoss: 28.584343\n",
      "Train Epoch: 30 [25600/60000 (43%)]\tLoss: 28.965866\n",
      "Train Epoch: 30 [38400/60000 (64%)]\tLoss: 28.085285\n",
      "Train Epoch: 30 [51200/60000 (85%)]\tLoss: 29.018261\n",
      "====> Epoch: 30 Average loss: 29.2246\n",
      "Train Epoch: 31 [0/60000 (0%)]\tLoss: 28.308573\n",
      "Train Epoch: 31 [12800/60000 (21%)]\tLoss: 28.426659\n",
      "Train Epoch: 31 [25600/60000 (43%)]\tLoss: 28.671312\n",
      "Train Epoch: 31 [38400/60000 (64%)]\tLoss: 28.769703\n",
      "Train Epoch: 31 [51200/60000 (85%)]\tLoss: 29.447746\n",
      "====> Epoch: 31 Average loss: 29.1902\n",
      "Train Epoch: 32 [0/60000 (0%)]\tLoss: 28.906204\n",
      "Train Epoch: 32 [12800/60000 (21%)]\tLoss: 28.802910\n",
      "Train Epoch: 32 [25600/60000 (43%)]\tLoss: 28.787796\n",
      "Train Epoch: 32 [38400/60000 (64%)]\tLoss: 29.242966\n",
      "Train Epoch: 32 [51200/60000 (85%)]\tLoss: 27.122261\n",
      "====> Epoch: 32 Average loss: 29.1671\n",
      "Train Epoch: 33 [0/60000 (0%)]\tLoss: 28.993528\n",
      "Train Epoch: 33 [12800/60000 (21%)]\tLoss: 28.808685\n",
      "Train Epoch: 33 [25600/60000 (43%)]\tLoss: 29.303862\n",
      "Train Epoch: 33 [38400/60000 (64%)]\tLoss: 30.223684\n",
      "Train Epoch: 33 [51200/60000 (85%)]\tLoss: 28.430973\n",
      "====> Epoch: 33 Average loss: 29.1397\n",
      "Train Epoch: 34 [0/60000 (0%)]\tLoss: 29.244783\n",
      "Train Epoch: 34 [12800/60000 (21%)]\tLoss: 28.805962\n",
      "Train Epoch: 34 [25600/60000 (43%)]\tLoss: 28.520126\n",
      "Train Epoch: 34 [38400/60000 (64%)]\tLoss: 28.836739\n",
      "Train Epoch: 34 [51200/60000 (85%)]\tLoss: 28.297935\n",
      "====> Epoch: 34 Average loss: 29.1157\n",
      "Train Epoch: 35 [0/60000 (0%)]\tLoss: 28.896467\n",
      "Train Epoch: 35 [12800/60000 (21%)]\tLoss: 28.600441\n",
      "Train Epoch: 35 [25600/60000 (43%)]\tLoss: 30.164991\n",
      "Train Epoch: 35 [38400/60000 (64%)]\tLoss: 29.248768\n",
      "Train Epoch: 35 [51200/60000 (85%)]\tLoss: 29.462864\n",
      "====> Epoch: 35 Average loss: 29.0598\n",
      "Train Epoch: 36 [0/60000 (0%)]\tLoss: 29.017696\n",
      "Train Epoch: 36 [12800/60000 (21%)]\tLoss: 29.893654\n",
      "Train Epoch: 36 [25600/60000 (43%)]\tLoss: 28.371273\n",
      "Train Epoch: 36 [38400/60000 (64%)]\tLoss: 29.273983\n",
      "Train Epoch: 36 [51200/60000 (85%)]\tLoss: 30.863380\n",
      "====> Epoch: 36 Average loss: 29.0621\n",
      "Train Epoch: 37 [0/60000 (0%)]\tLoss: 28.246193\n",
      "Train Epoch: 37 [12800/60000 (21%)]\tLoss: 28.988415\n",
      "Train Epoch: 37 [25600/60000 (43%)]\tLoss: 28.841141\n",
      "Train Epoch: 37 [38400/60000 (64%)]\tLoss: 29.597027\n",
      "Train Epoch: 37 [51200/60000 (85%)]\tLoss: 29.067900\n",
      "====> Epoch: 37 Average loss: 28.9968\n",
      "Train Epoch: 38 [0/60000 (0%)]\tLoss: 29.307808\n",
      "Train Epoch: 38 [12800/60000 (21%)]\tLoss: 28.939329\n",
      "Train Epoch: 38 [25600/60000 (43%)]\tLoss: 28.401358\n",
      "Train Epoch: 38 [38400/60000 (64%)]\tLoss: 29.553953\n",
      "Train Epoch: 38 [51200/60000 (85%)]\tLoss: 28.272371\n",
      "====> Epoch: 38 Average loss: 29.0059\n",
      "Train Epoch: 39 [0/60000 (0%)]\tLoss: 30.391144\n",
      "Train Epoch: 39 [12800/60000 (21%)]\tLoss: 28.406818\n",
      "Train Epoch: 39 [25600/60000 (43%)]\tLoss: 30.005447\n",
      "Train Epoch: 39 [38400/60000 (64%)]\tLoss: 29.470516\n",
      "Train Epoch: 39 [51200/60000 (85%)]\tLoss: 29.420704\n",
      "====> Epoch: 39 Average loss: 28.9817\n",
      "Train Epoch: 40 [0/60000 (0%)]\tLoss: 28.702423\n",
      "Train Epoch: 40 [12800/60000 (21%)]\tLoss: 29.854595\n",
      "Train Epoch: 40 [25600/60000 (43%)]\tLoss: 29.091972\n",
      "Train Epoch: 40 [38400/60000 (64%)]\tLoss: 29.304235\n",
      "Train Epoch: 40 [51200/60000 (85%)]\tLoss: 29.560081\n",
      "====> Epoch: 40 Average loss: 28.9468\n",
      "Train Epoch: 41 [0/60000 (0%)]\tLoss: 30.046064\n",
      "Train Epoch: 41 [12800/60000 (21%)]\tLoss: 29.413345\n",
      "Train Epoch: 41 [25600/60000 (43%)]\tLoss: 29.193851\n",
      "Train Epoch: 41 [38400/60000 (64%)]\tLoss: 29.791140\n",
      "Train Epoch: 41 [51200/60000 (85%)]\tLoss: 28.035534\n",
      "====> Epoch: 41 Average loss: 28.9054\n",
      "Train Epoch: 42 [0/60000 (0%)]\tLoss: 28.196026\n",
      "Train Epoch: 42 [12800/60000 (21%)]\tLoss: 28.140282\n",
      "Train Epoch: 42 [25600/60000 (43%)]\tLoss: 30.923813\n",
      "Train Epoch: 42 [38400/60000 (64%)]\tLoss: 29.160858\n",
      "Train Epoch: 42 [51200/60000 (85%)]\tLoss: 28.335207\n",
      "====> Epoch: 42 Average loss: 28.8936\n",
      "Train Epoch: 43 [0/60000 (0%)]\tLoss: 29.348726\n",
      "Train Epoch: 43 [12800/60000 (21%)]\tLoss: 29.724979\n",
      "Train Epoch: 43 [25600/60000 (43%)]\tLoss: 28.795273\n",
      "Train Epoch: 43 [38400/60000 (64%)]\tLoss: 28.607800\n",
      "Train Epoch: 43 [51200/60000 (85%)]\tLoss: 29.097221\n",
      "====> Epoch: 43 Average loss: 28.8623\n",
      "Train Epoch: 44 [0/60000 (0%)]\tLoss: 27.740559\n",
      "Train Epoch: 44 [12800/60000 (21%)]\tLoss: 28.217457\n",
      "Train Epoch: 44 [25600/60000 (43%)]\tLoss: 30.110420\n",
      "Train Epoch: 44 [38400/60000 (64%)]\tLoss: 28.716816\n",
      "Train Epoch: 44 [51200/60000 (85%)]\tLoss: 29.465700\n",
      "====> Epoch: 44 Average loss: 28.8530\n",
      "Train Epoch: 45 [0/60000 (0%)]\tLoss: 28.832981\n",
      "Train Epoch: 45 [12800/60000 (21%)]\tLoss: 28.763260\n",
      "Train Epoch: 45 [25600/60000 (43%)]\tLoss: 28.698666\n",
      "Train Epoch: 45 [38400/60000 (64%)]\tLoss: 29.087608\n",
      "Train Epoch: 45 [51200/60000 (85%)]\tLoss: 28.813656\n",
      "====> Epoch: 45 Average loss: 28.8436\n",
      "Train Epoch: 46 [0/60000 (0%)]\tLoss: 29.593136\n",
      "Train Epoch: 46 [12800/60000 (21%)]\tLoss: 28.942062\n",
      "Train Epoch: 46 [25600/60000 (43%)]\tLoss: 29.777401\n",
      "Train Epoch: 46 [38400/60000 (64%)]\tLoss: 27.552685\n",
      "Train Epoch: 46 [51200/60000 (85%)]\tLoss: 29.124855\n",
      "====> Epoch: 46 Average loss: 28.8081\n",
      "Train Epoch: 47 [0/60000 (0%)]\tLoss: 29.092293\n",
      "Train Epoch: 47 [12800/60000 (21%)]\tLoss: 28.847689\n",
      "Train Epoch: 47 [25600/60000 (43%)]\tLoss: 28.947968\n",
      "Train Epoch: 47 [38400/60000 (64%)]\tLoss: 29.061914\n",
      "Train Epoch: 47 [51200/60000 (85%)]\tLoss: 28.651283\n",
      "====> Epoch: 47 Average loss: 28.8098\n",
      "Train Epoch: 48 [0/60000 (0%)]\tLoss: 26.750555\n",
      "Train Epoch: 48 [12800/60000 (21%)]\tLoss: 28.219540\n",
      "Train Epoch: 48 [25600/60000 (43%)]\tLoss: 29.529823\n",
      "Train Epoch: 48 [38400/60000 (64%)]\tLoss: 29.518871\n",
      "Train Epoch: 48 [51200/60000 (85%)]\tLoss: 28.023201\n",
      "====> Epoch: 48 Average loss: 28.8097\n",
      "Train Epoch: 49 [0/60000 (0%)]\tLoss: 29.518541\n",
      "Train Epoch: 49 [12800/60000 (21%)]\tLoss: 28.782057\n",
      "Train Epoch: 49 [25600/60000 (43%)]\tLoss: 28.698990\n",
      "Train Epoch: 49 [38400/60000 (64%)]\tLoss: 29.736134\n",
      "Train Epoch: 49 [51200/60000 (85%)]\tLoss: 29.592495\n",
      "====> Epoch: 49 Average loss: 28.7625\n",
      "Train Epoch: 50 [0/60000 (0%)]\tLoss: 28.488533\n",
      "Train Epoch: 50 [12800/60000 (21%)]\tLoss: 29.548220\n",
      "Train Epoch: 50 [25600/60000 (43%)]\tLoss: 27.568245\n",
      "Train Epoch: 50 [38400/60000 (64%)]\tLoss: 28.363894\n",
      "Train Epoch: 50 [51200/60000 (85%)]\tLoss: 28.181688\n",
      "====> Epoch: 50 Average loss: 28.7454\n",
      "Train Epoch: 51 [0/60000 (0%)]\tLoss: 28.386768\n",
      "Train Epoch: 51 [12800/60000 (21%)]\tLoss: 29.105976\n",
      "Train Epoch: 51 [25600/60000 (43%)]\tLoss: 28.851517\n",
      "Train Epoch: 51 [38400/60000 (64%)]\tLoss: 29.006519\n",
      "Train Epoch: 51 [51200/60000 (85%)]\tLoss: 29.059067\n",
      "====> Epoch: 51 Average loss: 28.7202\n",
      "Train Epoch: 52 [0/60000 (0%)]\tLoss: 28.981844\n",
      "Train Epoch: 52 [12800/60000 (21%)]\tLoss: 30.064030\n",
      "Train Epoch: 52 [25600/60000 (43%)]\tLoss: 27.684519\n",
      "Train Epoch: 52 [38400/60000 (64%)]\tLoss: 28.846142\n",
      "Train Epoch: 52 [51200/60000 (85%)]\tLoss: 27.338268\n",
      "====> Epoch: 52 Average loss: 28.6937\n",
      "Train Epoch: 53 [0/60000 (0%)]\tLoss: 28.355881\n",
      "Train Epoch: 53 [12800/60000 (21%)]\tLoss: 25.701286\n",
      "Train Epoch: 53 [25600/60000 (43%)]\tLoss: 28.607300\n",
      "Train Epoch: 53 [38400/60000 (64%)]\tLoss: 29.345358\n",
      "Train Epoch: 53 [51200/60000 (85%)]\tLoss: 28.983433\n",
      "====> Epoch: 53 Average loss: 28.6735\n",
      "Train Epoch: 54 [0/60000 (0%)]\tLoss: 28.242289\n",
      "Train Epoch: 54 [12800/60000 (21%)]\tLoss: 28.109365\n",
      "Train Epoch: 54 [25600/60000 (43%)]\tLoss: 28.669321\n",
      "Train Epoch: 54 [38400/60000 (64%)]\tLoss: 28.637997\n",
      "Train Epoch: 54 [51200/60000 (85%)]\tLoss: 30.183723\n",
      "====> Epoch: 54 Average loss: 28.6820\n",
      "Train Epoch: 55 [0/60000 (0%)]\tLoss: 29.286980\n",
      "Train Epoch: 55 [12800/60000 (21%)]\tLoss: 28.416870\n",
      "Train Epoch: 55 [25600/60000 (43%)]\tLoss: 28.639240\n",
      "Train Epoch: 55 [38400/60000 (64%)]\tLoss: 28.319960\n",
      "Train Epoch: 55 [51200/60000 (85%)]\tLoss: 28.157757\n",
      "====> Epoch: 55 Average loss: 28.6385\n",
      "Train Epoch: 56 [0/60000 (0%)]\tLoss: 29.270454\n",
      "Train Epoch: 56 [12800/60000 (21%)]\tLoss: 29.999918\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Epoch: 56 [25600/60000 (43%)]\tLoss: 26.886427\n",
      "Train Epoch: 56 [38400/60000 (64%)]\tLoss: 28.099598\n",
      "Train Epoch: 56 [51200/60000 (85%)]\tLoss: 29.323936\n",
      "====> Epoch: 56 Average loss: 28.6220\n",
      "Train Epoch: 57 [0/60000 (0%)]\tLoss: 28.900547\n",
      "Train Epoch: 57 [12800/60000 (21%)]\tLoss: 27.997540\n",
      "Train Epoch: 57 [25600/60000 (43%)]\tLoss: 27.172966\n",
      "Train Epoch: 57 [38400/60000 (64%)]\tLoss: 28.231377\n",
      "Train Epoch: 57 [51200/60000 (85%)]\tLoss: 28.572067\n",
      "====> Epoch: 57 Average loss: 28.6236\n",
      "Train Epoch: 58 [0/60000 (0%)]\tLoss: 28.787979\n",
      "Train Epoch: 58 [12800/60000 (21%)]\tLoss: 28.486635\n",
      "Train Epoch: 58 [25600/60000 (43%)]\tLoss: 28.211384\n",
      "Train Epoch: 58 [38400/60000 (64%)]\tLoss: 28.844704\n",
      "Train Epoch: 58 [51200/60000 (85%)]\tLoss: 29.157852\n",
      "====> Epoch: 58 Average loss: 28.6225\n",
      "Train Epoch: 59 [0/60000 (0%)]\tLoss: 29.493900\n",
      "Train Epoch: 59 [12800/60000 (21%)]\tLoss: 28.969484\n",
      "Train Epoch: 59 [25600/60000 (43%)]\tLoss: 29.101681\n",
      "Train Epoch: 59 [38400/60000 (64%)]\tLoss: 29.139744\n",
      "Train Epoch: 59 [51200/60000 (85%)]\tLoss: 28.762054\n",
      "====> Epoch: 59 Average loss: 28.6247\n",
      "Train Epoch: 60 [0/60000 (0%)]\tLoss: 28.186562\n",
      "Train Epoch: 60 [12800/60000 (21%)]\tLoss: 28.305017\n",
      "Train Epoch: 60 [25600/60000 (43%)]\tLoss: 28.615343\n",
      "Train Epoch: 60 [38400/60000 (64%)]\tLoss: 28.206009\n",
      "Train Epoch: 60 [51200/60000 (85%)]\tLoss: 27.875103\n",
      "====> Epoch: 60 Average loss: 28.6068\n",
      "Train Epoch: 61 [0/60000 (0%)]\tLoss: 28.084419\n",
      "Train Epoch: 61 [12800/60000 (21%)]\tLoss: 27.766363\n",
      "Train Epoch: 61 [25600/60000 (43%)]\tLoss: 28.717920\n",
      "Train Epoch: 61 [38400/60000 (64%)]\tLoss: 29.540506\n",
      "Train Epoch: 61 [51200/60000 (85%)]\tLoss: 29.536980\n",
      "====> Epoch: 61 Average loss: 28.5736\n",
      "Train Epoch: 62 [0/60000 (0%)]\tLoss: 30.260386\n",
      "Train Epoch: 62 [12800/60000 (21%)]\tLoss: 28.702366\n",
      "Train Epoch: 62 [25600/60000 (43%)]\tLoss: 30.311745\n",
      "Train Epoch: 62 [38400/60000 (64%)]\tLoss: 29.303967\n",
      "Train Epoch: 62 [51200/60000 (85%)]\tLoss: 27.731878\n",
      "====> Epoch: 62 Average loss: 28.5730\n",
      "Train Epoch: 63 [0/60000 (0%)]\tLoss: 28.106098\n",
      "Train Epoch: 63 [12800/60000 (21%)]\tLoss: 29.471039\n",
      "Train Epoch: 63 [25600/60000 (43%)]\tLoss: 28.485422\n",
      "Train Epoch: 63 [38400/60000 (64%)]\tLoss: 27.741465\n",
      "Train Epoch: 63 [51200/60000 (85%)]\tLoss: 28.442802\n",
      "====> Epoch: 63 Average loss: 28.5459\n",
      "Train Epoch: 64 [0/60000 (0%)]\tLoss: 27.724640\n",
      "Train Epoch: 64 [12800/60000 (21%)]\tLoss: 30.057240\n",
      "Train Epoch: 64 [25600/60000 (43%)]\tLoss: 28.949724\n",
      "Train Epoch: 64 [38400/60000 (64%)]\tLoss: 29.427559\n",
      "Train Epoch: 64 [51200/60000 (85%)]\tLoss: 28.377747\n",
      "====> Epoch: 64 Average loss: 28.5625\n",
      "Train Epoch: 65 [0/60000 (0%)]\tLoss: 28.785015\n",
      "Train Epoch: 65 [12800/60000 (21%)]\tLoss: 27.660938\n",
      "Train Epoch: 65 [25600/60000 (43%)]\tLoss: 28.587070\n",
      "Train Epoch: 65 [38400/60000 (64%)]\tLoss: 26.625290\n",
      "Train Epoch: 65 [51200/60000 (85%)]\tLoss: 27.357773\n",
      "====> Epoch: 65 Average loss: 28.5275\n",
      "Train Epoch: 66 [0/60000 (0%)]\tLoss: 28.450836\n",
      "Train Epoch: 66 [12800/60000 (21%)]\tLoss: 28.861435\n",
      "Train Epoch: 66 [25600/60000 (43%)]\tLoss: 29.998249\n",
      "Train Epoch: 66 [38400/60000 (64%)]\tLoss: 28.349022\n",
      "Train Epoch: 66 [51200/60000 (85%)]\tLoss: 28.922348\n",
      "====> Epoch: 66 Average loss: 28.5173\n",
      "Train Epoch: 67 [0/60000 (0%)]\tLoss: 29.511518\n",
      "Train Epoch: 67 [12800/60000 (21%)]\tLoss: 28.616154\n",
      "Train Epoch: 67 [25600/60000 (43%)]\tLoss: 28.316280\n",
      "Train Epoch: 67 [38400/60000 (64%)]\tLoss: 28.890785\n",
      "Train Epoch: 67 [51200/60000 (85%)]\tLoss: 29.365330\n",
      "====> Epoch: 67 Average loss: 28.5081\n",
      "Train Epoch: 68 [0/60000 (0%)]\tLoss: 28.279198\n",
      "Train Epoch: 68 [12800/60000 (21%)]\tLoss: 28.077579\n",
      "Train Epoch: 68 [25600/60000 (43%)]\tLoss: 28.694098\n",
      "Train Epoch: 68 [38400/60000 (64%)]\tLoss: 28.804489\n",
      "Train Epoch: 68 [51200/60000 (85%)]\tLoss: 27.846481\n",
      "====> Epoch: 68 Average loss: 28.4895\n",
      "Train Epoch: 69 [0/60000 (0%)]\tLoss: 28.607380\n",
      "Train Epoch: 69 [12800/60000 (21%)]\tLoss: 28.834291\n",
      "Train Epoch: 69 [25600/60000 (43%)]\tLoss: 28.623692\n",
      "Train Epoch: 69 [38400/60000 (64%)]\tLoss: 27.871918\n",
      "Train Epoch: 69 [51200/60000 (85%)]\tLoss: 30.008867\n",
      "====> Epoch: 69 Average loss: 28.4911\n",
      "Train Epoch: 70 [0/60000 (0%)]\tLoss: 28.181942\n",
      "Train Epoch: 70 [12800/60000 (21%)]\tLoss: 29.545860\n",
      "Train Epoch: 70 [25600/60000 (43%)]\tLoss: 29.403370\n",
      "Train Epoch: 70 [38400/60000 (64%)]\tLoss: 28.258066\n",
      "Train Epoch: 70 [51200/60000 (85%)]\tLoss: 27.643871\n",
      "====> Epoch: 70 Average loss: 28.4675\n",
      "Train Epoch: 71 [0/60000 (0%)]\tLoss: 27.957352\n",
      "Train Epoch: 71 [12800/60000 (21%)]\tLoss: 29.409380\n",
      "Train Epoch: 71 [25600/60000 (43%)]\tLoss: 27.930595\n",
      "Train Epoch: 71 [38400/60000 (64%)]\tLoss: 28.761438\n",
      "Train Epoch: 71 [51200/60000 (85%)]\tLoss: 28.926760\n",
      "====> Epoch: 71 Average loss: 28.4674\n",
      "Train Epoch: 72 [0/60000 (0%)]\tLoss: 28.314390\n",
      "Train Epoch: 72 [12800/60000 (21%)]\tLoss: 28.184330\n",
      "Train Epoch: 72 [25600/60000 (43%)]\tLoss: 27.651810\n",
      "Train Epoch: 72 [38400/60000 (64%)]\tLoss: 27.820982\n",
      "Train Epoch: 72 [51200/60000 (85%)]\tLoss: 29.494406\n",
      "====> Epoch: 72 Average loss: 28.4578\n",
      "Train Epoch: 73 [0/60000 (0%)]\tLoss: 28.052242\n",
      "Train Epoch: 73 [12800/60000 (21%)]\tLoss: 29.789310\n",
      "Train Epoch: 73 [25600/60000 (43%)]\tLoss: 26.808613\n",
      "Train Epoch: 73 [38400/60000 (64%)]\tLoss: 27.801287\n",
      "Train Epoch: 73 [51200/60000 (85%)]\tLoss: 28.819286\n",
      "====> Epoch: 73 Average loss: 28.4468\n",
      "Train Epoch: 74 [0/60000 (0%)]\tLoss: 28.896933\n",
      "Train Epoch: 74 [12800/60000 (21%)]\tLoss: 28.850624\n",
      "Train Epoch: 74 [25600/60000 (43%)]\tLoss: 29.054195\n",
      "Train Epoch: 74 [38400/60000 (64%)]\tLoss: 30.144028\n",
      "Train Epoch: 74 [51200/60000 (85%)]\tLoss: 28.931324\n",
      "====> Epoch: 74 Average loss: 28.4298\n",
      "Train Epoch: 75 [0/60000 (0%)]\tLoss: 27.949577\n",
      "Train Epoch: 75 [12800/60000 (21%)]\tLoss: 29.712132\n",
      "Train Epoch: 75 [25600/60000 (43%)]\tLoss: 28.787903\n",
      "Train Epoch: 75 [38400/60000 (64%)]\tLoss: 28.030304\n",
      "Train Epoch: 75 [51200/60000 (85%)]\tLoss: 28.312954\n",
      "====> Epoch: 75 Average loss: 28.3885\n",
      "Train Epoch: 76 [0/60000 (0%)]\tLoss: 29.296371\n",
      "Train Epoch: 76 [12800/60000 (21%)]\tLoss: 29.144676\n",
      "Train Epoch: 76 [25600/60000 (43%)]\tLoss: 29.963020\n",
      "Train Epoch: 76 [38400/60000 (64%)]\tLoss: 27.943571\n",
      "Train Epoch: 76 [51200/60000 (85%)]\tLoss: 29.240404\n",
      "====> Epoch: 76 Average loss: 28.4267\n",
      "Train Epoch: 77 [0/60000 (0%)]\tLoss: 27.747213\n",
      "Train Epoch: 77 [12800/60000 (21%)]\tLoss: 27.684841\n",
      "Train Epoch: 77 [25600/60000 (43%)]\tLoss: 27.523190\n",
      "Train Epoch: 77 [38400/60000 (64%)]\tLoss: 29.017170\n",
      "Train Epoch: 77 [51200/60000 (85%)]\tLoss: 27.871475\n",
      "====> Epoch: 77 Average loss: 28.4115\n",
      "Train Epoch: 78 [0/60000 (0%)]\tLoss: 26.686985\n",
      "Train Epoch: 78 [12800/60000 (21%)]\tLoss: 27.858383\n",
      "Train Epoch: 78 [25600/60000 (43%)]\tLoss: 28.465462\n",
      "Train Epoch: 78 [38400/60000 (64%)]\tLoss: 27.476528\n",
      "Train Epoch: 78 [51200/60000 (85%)]\tLoss: 29.315985\n",
      "====> Epoch: 78 Average loss: 28.3759\n",
      "Train Epoch: 79 [0/60000 (0%)]\tLoss: 27.827763\n",
      "Train Epoch: 79 [12800/60000 (21%)]\tLoss: 27.419086\n",
      "Train Epoch: 79 [25600/60000 (43%)]\tLoss: 30.189674\n",
      "Train Epoch: 79 [38400/60000 (64%)]\tLoss: 28.323854\n",
      "Train Epoch: 79 [51200/60000 (85%)]\tLoss: 28.649834\n",
      "====> Epoch: 79 Average loss: 28.3616\n",
      "Train Epoch: 80 [0/60000 (0%)]\tLoss: 26.943005\n",
      "Train Epoch: 80 [12800/60000 (21%)]\tLoss: 28.206684\n",
      "Train Epoch: 80 [25600/60000 (43%)]\tLoss: 28.294369\n",
      "Train Epoch: 80 [38400/60000 (64%)]\tLoss: 27.483158\n",
      "Train Epoch: 80 [51200/60000 (85%)]\tLoss: 28.808107\n",
      "====> Epoch: 80 Average loss: 28.3766\n",
      "Train Epoch: 81 [0/60000 (0%)]\tLoss: 27.441273\n",
      "Train Epoch: 81 [12800/60000 (21%)]\tLoss: 27.350582\n",
      "Train Epoch: 81 [25600/60000 (43%)]\tLoss: 28.023540\n",
      "Train Epoch: 81 [38400/60000 (64%)]\tLoss: 29.376869\n",
      "Train Epoch: 81 [51200/60000 (85%)]\tLoss: 28.101452\n",
      "====> Epoch: 81 Average loss: 28.3579\n",
      "Train Epoch: 82 [0/60000 (0%)]\tLoss: 28.451614\n",
      "Train Epoch: 82 [12800/60000 (21%)]\tLoss: 28.089062\n",
      "Train Epoch: 82 [25600/60000 (43%)]\tLoss: 29.289324\n",
      "Train Epoch: 82 [38400/60000 (64%)]\tLoss: 29.265394\n",
      "Train Epoch: 82 [51200/60000 (85%)]\tLoss: 29.066242\n",
      "====> Epoch: 82 Average loss: 28.3684\n",
      "Train Epoch: 83 [0/60000 (0%)]\tLoss: 27.552654\n",
      "Train Epoch: 83 [12800/60000 (21%)]\tLoss: 27.917965\n",
      "Train Epoch: 83 [25600/60000 (43%)]\tLoss: 28.593105\n",
      "Train Epoch: 83 [38400/60000 (64%)]\tLoss: 29.037209\n",
      "Train Epoch: 83 [51200/60000 (85%)]\tLoss: 28.034590\n",
      "====> Epoch: 83 Average loss: 28.3689\n",
      "Train Epoch: 84 [0/60000 (0%)]\tLoss: 29.631090\n",
      "Train Epoch: 84 [12800/60000 (21%)]\tLoss: 27.932201\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Epoch: 84 [25600/60000 (43%)]\tLoss: 27.612846\n",
      "Train Epoch: 84 [38400/60000 (64%)]\tLoss: 27.958817\n",
      "Train Epoch: 84 [51200/60000 (85%)]\tLoss: 28.648518\n",
      "====> Epoch: 84 Average loss: 28.3492\n",
      "Train Epoch: 85 [0/60000 (0%)]\tLoss: 29.203257\n",
      "Train Epoch: 85 [12800/60000 (21%)]\tLoss: 28.752777\n",
      "Train Epoch: 85 [25600/60000 (43%)]\tLoss: 29.391205\n",
      "Train Epoch: 85 [38400/60000 (64%)]\tLoss: 28.371208\n",
      "Train Epoch: 85 [51200/60000 (85%)]\tLoss: 27.343235\n",
      "====> Epoch: 85 Average loss: 28.3183\n",
      "Train Epoch: 86 [0/60000 (0%)]\tLoss: 28.250309\n",
      "Train Epoch: 86 [12800/60000 (21%)]\tLoss: 28.119259\n",
      "Train Epoch: 86 [25600/60000 (43%)]\tLoss: 28.935797\n",
      "Train Epoch: 86 [38400/60000 (64%)]\tLoss: 28.092178\n",
      "Train Epoch: 86 [51200/60000 (85%)]\tLoss: 28.010101\n",
      "====> Epoch: 86 Average loss: 28.3190\n",
      "Train Epoch: 87 [0/60000 (0%)]\tLoss: 29.455128\n",
      "Train Epoch: 87 [12800/60000 (21%)]\tLoss: 28.264370\n",
      "Train Epoch: 87 [25600/60000 (43%)]\tLoss: 28.342337\n",
      "Train Epoch: 87 [38400/60000 (64%)]\tLoss: 27.640537\n",
      "Train Epoch: 87 [51200/60000 (85%)]\tLoss: 27.257610\n",
      "====> Epoch: 87 Average loss: 28.3396\n",
      "Train Epoch: 88 [0/60000 (0%)]\tLoss: 27.681135\n",
      "Train Epoch: 88 [12800/60000 (21%)]\tLoss: 28.313272\n",
      "Train Epoch: 88 [25600/60000 (43%)]\tLoss: 26.566006\n",
      "Train Epoch: 88 [38400/60000 (64%)]\tLoss: 29.389267\n",
      "Train Epoch: 88 [51200/60000 (85%)]\tLoss: 29.921810\n",
      "====> Epoch: 88 Average loss: 28.3108\n",
      "Train Epoch: 89 [0/60000 (0%)]\tLoss: 27.515110\n",
      "Train Epoch: 89 [12800/60000 (21%)]\tLoss: 28.540718\n",
      "Train Epoch: 89 [25600/60000 (43%)]\tLoss: 28.845592\n",
      "Train Epoch: 89 [38400/60000 (64%)]\tLoss: 29.146767\n",
      "Train Epoch: 89 [51200/60000 (85%)]\tLoss: 28.398174\n",
      "====> Epoch: 89 Average loss: 28.3044\n",
      "Train Epoch: 90 [0/60000 (0%)]\tLoss: 27.768028\n",
      "Train Epoch: 90 [12800/60000 (21%)]\tLoss: 28.987511\n",
      "Train Epoch: 90 [25600/60000 (43%)]\tLoss: 29.068241\n",
      "Train Epoch: 90 [38400/60000 (64%)]\tLoss: 29.129974\n",
      "Train Epoch: 90 [51200/60000 (85%)]\tLoss: 26.976212\n",
      "====> Epoch: 90 Average loss: 28.2984\n",
      "Train Epoch: 91 [0/60000 (0%)]\tLoss: 28.668402\n",
      "Train Epoch: 91 [12800/60000 (21%)]\tLoss: 27.478672\n",
      "Train Epoch: 91 [25600/60000 (43%)]\tLoss: 27.967779\n",
      "Train Epoch: 91 [38400/60000 (64%)]\tLoss: 27.689548\n",
      "Train Epoch: 91 [51200/60000 (85%)]\tLoss: 27.981226\n",
      "====> Epoch: 91 Average loss: 28.3028\n",
      "Train Epoch: 92 [0/60000 (0%)]\tLoss: 28.172039\n",
      "Train Epoch: 92 [12800/60000 (21%)]\tLoss: 27.905527\n",
      "Train Epoch: 92 [25600/60000 (43%)]\tLoss: 28.349701\n",
      "Train Epoch: 92 [38400/60000 (64%)]\tLoss: 28.185818\n",
      "Train Epoch: 92 [51200/60000 (85%)]\tLoss: 28.006760\n",
      "====> Epoch: 92 Average loss: 28.2962\n",
      "Train Epoch: 93 [0/60000 (0%)]\tLoss: 28.137047\n",
      "Train Epoch: 93 [12800/60000 (21%)]\tLoss: 27.461464\n",
      "Train Epoch: 93 [25600/60000 (43%)]\tLoss: 27.367018\n",
      "Train Epoch: 93 [38400/60000 (64%)]\tLoss: 28.600716\n",
      "Train Epoch: 93 [51200/60000 (85%)]\tLoss: 26.852875\n",
      "====> Epoch: 93 Average loss: 28.2586\n",
      "Train Epoch: 94 [0/60000 (0%)]\tLoss: 28.658562\n",
      "Train Epoch: 94 [12800/60000 (21%)]\tLoss: 27.871185\n",
      "Train Epoch: 94 [25600/60000 (43%)]\tLoss: 29.044031\n",
      "Train Epoch: 94 [38400/60000 (64%)]\tLoss: 29.269758\n",
      "Train Epoch: 94 [51200/60000 (85%)]\tLoss: 28.659653\n",
      "====> Epoch: 94 Average loss: 28.2643\n",
      "Train Epoch: 95 [0/60000 (0%)]\tLoss: 27.919479\n",
      "Train Epoch: 95 [12800/60000 (21%)]\tLoss: 27.762722\n",
      "Train Epoch: 95 [25600/60000 (43%)]\tLoss: 26.691334\n",
      "Train Epoch: 95 [38400/60000 (64%)]\tLoss: 28.898460\n",
      "Train Epoch: 95 [51200/60000 (85%)]\tLoss: 28.239012\n",
      "====> Epoch: 95 Average loss: 28.2579\n",
      "Train Epoch: 96 [0/60000 (0%)]\tLoss: 26.803535\n",
      "Train Epoch: 96 [12800/60000 (21%)]\tLoss: 27.288227\n",
      "Train Epoch: 96 [25600/60000 (43%)]\tLoss: 27.901173\n",
      "Train Epoch: 96 [38400/60000 (64%)]\tLoss: 29.518402\n",
      "Train Epoch: 96 [51200/60000 (85%)]\tLoss: 29.076067\n",
      "====> Epoch: 96 Average loss: 28.2582\n",
      "Train Epoch: 97 [0/60000 (0%)]\tLoss: 29.066921\n",
      "Train Epoch: 97 [12800/60000 (21%)]\tLoss: 28.440113\n",
      "Train Epoch: 97 [25600/60000 (43%)]\tLoss: 26.841232\n",
      "Train Epoch: 97 [38400/60000 (64%)]\tLoss: 26.914097\n",
      "Train Epoch: 97 [51200/60000 (85%)]\tLoss: 28.268524\n",
      "====> Epoch: 97 Average loss: 28.2478\n",
      "Train Epoch: 98 [0/60000 (0%)]\tLoss: 29.353155\n",
      "Train Epoch: 98 [12800/60000 (21%)]\tLoss: 28.510120\n",
      "Train Epoch: 98 [25600/60000 (43%)]\tLoss: 29.327541\n",
      "Train Epoch: 98 [38400/60000 (64%)]\tLoss: 28.147558\n",
      "Train Epoch: 98 [51200/60000 (85%)]\tLoss: 29.402172\n",
      "====> Epoch: 98 Average loss: 28.2285\n",
      "Train Epoch: 99 [0/60000 (0%)]\tLoss: 26.748432\n",
      "Train Epoch: 99 [12800/60000 (21%)]\tLoss: 27.822241\n",
      "Train Epoch: 99 [25600/60000 (43%)]\tLoss: 27.896614\n",
      "Train Epoch: 99 [38400/60000 (64%)]\tLoss: 29.203022\n",
      "Train Epoch: 99 [51200/60000 (85%)]\tLoss: 28.939926\n",
      "====> Epoch: 99 Average loss: 28.2367\n"
     ]
    }
   ],
   "source": [
    "num_epochs = 100\n",
    "batch_size = 128\n",
    "learning_rate = 1e-3\n",
    "\n",
    "img_transform = transforms.Compose([\n",
    "    transforms.ToTensor()\n",
    "    # transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n",
    "])\n",
    "\n",
    "dataset = MNIST('./data', transform=img_transform, download=True)\n",
    "dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)\n",
    "\n",
    "model = VAE()\n",
    "if torch.cuda.is_available():\n",
    "    model.cuda()\n",
    "\n",
    "reconstruction_function = nn.MSELoss(size_average=False)\n",
    "\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-3)\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    train_loss = 0\n",
    "    for batch_idx, data in enumerate(dataloader):\n",
    "        img, _ = data\n",
    "        img = img.view(img.size(0), -1)\n",
    "        img = Variable(img)\n",
    "        if torch.cuda.is_available():\n",
    "            img = img.cuda()\n",
    "        optimizer.zero_grad()\n",
    "        recon_batch, mu, logvar = model(img)\n",
    "        loss = loss_function(recon_batch, img, mu, logvar)\n",
    "        loss.backward()\n",
    "        train_loss += loss.data[0]\n",
    "        optimizer.step()\n",
    "        if batch_idx % 100 == 0:\n",
    "            print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
    "                epoch,\n",
    "                batch_idx * len(img),\n",
    "                len(dataloader.dataset), 100. * batch_idx / len(dataloader),\n",
    "                loss.data[0] / len(img)))\n",
    "\n",
    "    print('====> Epoch: {} Average loss: {:.4f}'.format(\n",
    "        epoch, train_loss / len(dataloader.dataset)))\n",
    "    if epoch % 10 == 0:\n",
    "        save = to_img(recon_batch.cpu().data)\n",
    "        save_image(save, './vae_img/image_{}.png'.format(epoch))\n",
    "\n",
    "torch.save(model.state_dict(), './vae.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
